import sys
from datetime import datetime

from PySide2.QtWidgets import QApplication, QMainWindow, QFileDialog, QGraphicsScene, QGraphicsPixmapItem, QPushButton, \
    QMessageBox, QProgressDialog
from PySide2.QtCore import QThread, Signal
from PySide2.QtUiTools import QUiLoader
from PySide2.QtGui import QPixmap, Qt
from PySide2.QtGui import  QIcon
import shutil
import os
import train
import predict
import threading


class TrainThread(QThread):
    output_signal = Signal(str)
    pause_signal = Signal(bool)
    update_ui_signal = Signal(str, str, str, str)

    def __init__(self):
        super().__init__()
        self._pause_event = threading.Event()
        self._pause_event.set()
        self.pause_signal.connect(self.set_pause)
        self.update_ui_signal.connect(self.update_ui_callback)

    def run(self):
        import sys
        from io import StringIO

        class StreamRedirector(StringIO):
            def __init__(self, output_signal):
                super().__init__()
                self.output_signal = output_signal

            def write(self, s):
                super().write(s)
                self.output_signal.emit(s)

        # Redirect stdout
        old_stdout = sys.stdout
        sys.stdout = StreamRedirector(self.output_signal)

        try:
            args = train.parse_args()  # Get the args from train.py
            train.main(args, self._pause_event,
                       self.update_ui_callback)  # Pass the pause event and update callback to main
        finally:
            sys.stdout = old_stdout

    def set_pause(self, pause):
        if pause:
            self._pause_event.clear()
        else:
            self._pause_event.set()

    def update_ui_callback(self, source_folder1, old_filename, target_folder2, new_filename):
        self.update_ui_signal.emit(source_folder1, old_filename, target_folder2, new_filename)


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.initUI()

    def initUI(self):
        self.ui = QUiLoader().load('main.ui')
        self.setCentralWidget(self.ui)
        self.setWindowTitle("Medical Image Segmentation System ")

        self.ui.uploadfolder.clicked.connect(self.upload_folder)
        self.ui.begintrain.clicked.connect(self.start_training)
        self.ui.pauseButton.clicked.connect(self.pause_training)
        self.ui.uploadResultButton.clicked.connect(self.upload_segmentation_result)
        self.ui.downloadWeightsButton.clicked.connect(self.download_weights)
        self.ui.uploadImageButton.clicked.connect(self.upload_image)
        self.ui.uploadWeightsButton.clicked.connect(self.upload_weights)

        # Initialize QGraphicsView and QGraphicsScene
        self.left_scene = QGraphicsScene()
        self.ui.graphicsView_left.setScene(self.left_scene)

        self.right_scene = QGraphicsScene()
        self.ui.graphicsView_right.setScene(self.right_scene)

        self.left_scene_2 = QGraphicsScene()
        self.ui.graphicsView_left_2.setScene(self.left_scene_2)

        self.right_scene_2 = QGraphicsScene()
        self.ui.graphicsView_right_2.setScene(self.right_scene_2)

        self.source_folder1 = ""
        self.old_filename = ""
        self.target_folder2 = ""
        self.new_filename = ""

    def upload_folder(self):
        folder_path = QFileDialog.getExistingDirectory(self, "Select Folder")

        if folder_path:
            ui_file_path = os.path.abspath(__file__)
            ui_folder_path = os.path.dirname(ui_file_path)

            destination_folder = os.path.join(ui_folder_path, "SLIVER07")
            if os.path.exists(destination_folder):
                shutil.rmtree(destination_folder)
                print(f"Deleted existing folder: {destination_folder}")
            shutil.copytree(folder_path, destination_folder)
            print("Finish upload")

            # Show success message box
            msg_box = QMessageBox()
            msg_box.setIcon(QMessageBox.Information)
            msg_box.setText("文件上传成功")
            msg_box.setWindowTitle("Success")
            msg_box.setStandardButtons(QMessageBox.Ok)
            msg_box.exec_()
        else:
            # Show success message box
            msg_box = QMessageBox()
            msg_box.setIcon(QMessageBox.Information)
            msg_box.setText("文件上传失败")
            msg_box.setWindowTitle("False")
            msg_box.setStandardButtons(QMessageBox.Ok)
            msg_box.exec_()


    def start_training(self):
        self.thread = TrainThread()
        self.thread.output_signal.connect(self.update_text_browser)
        self.thread.update_ui_signal.connect(self.show_image_for_marking)
        self.thread.start()

    def pause_training(self):
        if self.thread._pause_event.is_set():
            self.thread.pause_signal.emit(True)
            self.ui.pauseButton.setText("恢复")
        else:
            self.thread.pause_signal.emit(False)
            self.ui.pauseButton.setText("暂停")

    def update_text_browser(self, text):
        self.ui.textBrowser.append(text)

    def show_image_for_marking(self, source_folder1, old_filename, target_folder2, new_filename):
        self.source_folder1 = source_folder1
        self.old_filename = old_filename
        self.target_folder2 = target_folder2
        self.new_filename = new_filename

        image_path = os.path.join(source_folder1, old_filename)
        pixmap = QPixmap(image_path)

        # Display image in QGraphicsView
        self.left_scene.clear()
        self.left_scene.addItem(QGraphicsPixmapItem(pixmap))

    def upload_segmentation_result(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Upload Segmentation Result")

        if file_path:
            dest_path = os.path.join(self.target_folder2, self.new_filename)
            shutil.copyfile(file_path, dest_path)
            print(f"Uploaded segmentation result to: {dest_path}")

            # Show the uploaded image
            pixmap = QPixmap(dest_path)

            # Display uploaded image in QGraphicsView
            self.right_scene.clear()
            self.right_scene.addItem(QGraphicsPixmapItem(pixmap))

            # Resume the training
            self.thread.pause_signal.emit(False)

    def download_weights(self):
        weights_file = os.path.join(os.path.dirname(__file__), 'save_weights', 'best_model.pth')
        if os.path.exists(weights_file):
            save_path, _ = QFileDialog.getSaveFileName(self, "Save Weights", "", "PyTorch Model (*.pth)")
            if save_path:
                shutil.copyfile(weights_file, save_path)
                print(f"Downloaded weights to: {save_path}")

                # Show success message box
                msg_box = QMessageBox()
                msg_box.setIcon(QMessageBox.Information)
                msg_box.setText(f"成功下载权重到 {save_path}")
                msg_box.setWindowTitle("Success")
                msg_box.setStandardButtons(QMessageBox.Ok)
                msg_box.exec_()
        else:
            print("Weights file does not exist.")

            # Show success message box
            msg_box = QMessageBox()
            msg_box.setIcon(QMessageBox.Information)
            msg_box.setText("权重文件不存在")
            msg_box.setWindowTitle("False")
            msg_box.setStandardButtons(QMessageBox.Ok)
            msg_box.exec_()

    def upload_image(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select Image")

        if file_path:
            # Show progress dialog
            self.progress_dialog = QProgressDialog("Processing image...", "Cancel", 0, 100, self)
            self.progress_dialog.setWindowTitle("Progress")
            self.progress_dialog.setWindowModality(Qt.WindowModal)
            self.progress_dialog.setValue(0)
            # Set fixed size for the progress dialog
            self.progress_dialog.setFixedSize(400, 100)
            self.progress_dialog.show()

            # Generate new filename based on current time
            new_filename = datetime.now().strftime("%Y%m%d%H%M%S") + os.path.splitext(file_path)[1]

            # Define the destination folder
            ui_file_path = os.path.abspath(__file__)
            ui_folder_path = os.path.dirname(ui_file_path)
            destination_folder = os.path.join(ui_folder_path, "predictdir", "vol")
            if not os.path.exists(destination_folder):
                os.makedirs(destination_folder)

            # Copy file to destination with new filename
            dest_path = os.path.join(destination_folder, new_filename)
            shutil.copyfile(file_path, dest_path)
            print(f"Uploaded image to: {dest_path}")

            # Display image in QGraphicsView
            pixmap = QPixmap(dest_path)
            self.left_scene_2.clear()
            self.left_scene_2.addItem(QGraphicsPixmapItem(pixmap))

            # Start image processing in a separate thread to avoid blocking the UI
            self.thread = ImageProcessingThread(dest_path)
            self.thread.progress_signal.connect(self.update_progress)
            self.thread.finished_signal.connect(self.processing_finished)
            self.thread.start()

    def update_progress(self, value):
        self.progress_dialog.setValue(value)

    def processing_finished(self, segmentation_result_path):
        self.progress_dialog.setValue(100)
        self.progress_dialog.close()

        # Display the segmentation result in the right graphics view
        pixmap = QPixmap(segmentation_result_path)
        self.right_scene_2.clear()
        self.right_scene_2.addItem(QGraphicsPixmapItem(pixmap))

    def upload_weights(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select Weights File", "", "PyTorch Model (*.pth)")

        if file_path:
            # Define the destination path
            ui_file_path = os.path.abspath(__file__)
            ui_folder_path = os.path.dirname(ui_file_path)
            destination_folder = os.path.join(ui_folder_path, "save_weights")
            if not os.path.exists(destination_folder):
                os.makedirs(destination_folder)

            dest_path = os.path.join(destination_folder, "best_model.pth")
            shutil.copyfile(file_path, dest_path)
            print(f"Uploaded weights to: {dest_path}")

            # Show success message box
            msg_box = QMessageBox()
            msg_box.setIcon(QMessageBox.Information)
            msg_box.setText("权重加载成功")
            msg_box.setWindowTitle("Success")
            msg_box.setStandardButtons(QMessageBox.Ok)
            msg_box.exec_()
        else:
            # Show success message box
            msg_box = QMessageBox()
            msg_box.setIcon(QMessageBox.Information)
            msg_box.setText("权重加载失败")
            msg_box.setWindowTitle("False")
            msg_box.setStandardButtons(QMessageBox.Ok)
            msg_box.exec_()

class ImageProcessingThread(QThread):
        progress_signal = Signal(int)
        finished_signal = Signal(str)

        def __init__(self, image_path):
            super().__init__()
            self.image_path = image_path

        def run(self):
            # Call predict.py's main function with the image path
            segmentation_result_path = predict.main(self.image_path, self.progress_signal)
            self.finished_signal.emit(segmentation_result_path)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    # 加载 icon
    app.setWindowIcon(QIcon('logo.png'))
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
